from yolov5.train import run

if __name__ == '__main__':
    run(
        data='data.yaml',                  # 数据地址以及类型设置
        weights=False,                     # 或你的模型 pt 路径
        cfg='yolov5/models/yolov5s.yaml',  # 模型源码地址
        imgsz=640,
        batch=16,                          # 一次性训练的图片数量
        epochs=100,                        # 训练轮次
        name='docker_yolo_train',
        project='runs/train'
    )
